import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cmx
import numpy as np
import os, glob, sys, subprocess
from matplotlib.colors import LinearSegmentedColormap
from math import pi

norm = np.linalg.norm

matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'
matplotlib.rcParams['font.size'] = 10
tabcolors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple',
                       'tab:brown', 'tab:pink', 'tab:gray', 'tab:olive', 'tab:cyan']


def main():

    xmin = float(sys.argv[1])
    xmax = float(sys.argv[2])
    name = sys.argv[3]
    emin = 1000
    emax = -1000
    nbin = 100
    temp = 120

    alldata = []
    fig_xi, axs_xi = plt.subplots(figsize=(3.4,2.5), dpi=300)
    fig_angle, axs_angle = plt.subplots(1, 2, figsize=(6.8,2.5), dpi=300)
    fig_xi_e, axs_xi_e = plt.subplots(figsize=(3.4,2.5), dpi=300)
    fig_xi_e_all, axs_xi_e_all = plt.subplots(figsize=(3.4,2.5), dpi=300)
    fig_step_e_all, axs_step_e_all = plt.subplots(figsize=(3.4,2.5), dpi=300)
    fig_step_xi_all, axs_step_xi_all = plt.subplots(figsize=(3.4,2.5), dpi=300)
    fig_xi_angle, axs_xi_angle = plt.subplots(figsize=(3.4,2.5), dpi=300)
    fout2 = open("npzlist", "w+")
    with open(f"metadata", "w+") as fout:
        for zfile in glob.glob(f"umb*/*.npz"):
            data = np.load(zfile)
            splt = zfile.split("/")
            xi = data['xi']
            e = data['alle'][:, 1]
            if xi.shape[0]>0 and np.max(e)<5:

                intc = data['colvar']

                axs_xi.hist(xi, bins=50)
                axs_angle[0].scatter(intc[::20, 0], intc[::20, 1],
                        color=tabcolors[data['basin']])
                axs_angle[1].scatter(intc[::20, 0], intc[::20, 1],
                        c=xi[::20], vmin=xmin, vmax=xmax)
                axs_xi_e.scatter(xi[::20], e[::20])
                axs_xi_angle.scatter(xi[::20], intc[::20, 0], color='k')
                axs_xi_angle.scatter(xi[::20], intc[::20, 1], color='r')

                h, x = np.histogram(xi)
                print(zfile, data['basin'], data['ends'],
                      np.max(h), np.average(xi), np.std(xi),
                      np.max(xi)-np.min(xi), file=fout2)
                cemax = np.max(e)
                print(data['center'], data['k'], file=fout)
                d = {}
                head = 1000
                d['xi'] = xi[head:]
                d['colvar'] = intc[head:, 0:2]
                d['x0'] = data['center']
                d['k'] = data['k']
                if 'k0' in data:
                    d['k0'] = data['k0']
                else:
                    d['k0'] = 0.
                alldata += [d]
            else:
                print("skip", zfile, np.max(e), data['k'])
            axs_xi_e_all.scatter(xi[::20], e[::20])
            axs_step_e_all.plot(np.arange(len(e[::20])), e[::20])
            axs_step_xi_all.plot(np.arange(len(xi[::20])), xi[::20])
    np.savez("alldata_uwham.npz", alldata=alldata)
    print("save uwham.npz")

    fout2.close()
    fig_xi.savefig(f"{name}_hist.png", dpi=300)
    fig_angle.savefig(f"{name}_scatter2d.png", dpi=300)
    fig_xi_e.savefig(f"{name}_xi_pe.png", dpi=300)
    axs_xi_e_all.set_ylim([-20, 40])
    fig_xi_e_all.savefig(f"{name}_xi_pe_all.png", dpi=300)
    axs_step_e_all.set_ylim([-20, 40])
    fig_step_e_all.savefig(f"{name}_step_pe_all.png", dpi=300)
    fig_step_xi_all.savefig(f"{name}_step_xi_all.png", dpi=300)
    fig_xi_angle.savefig(f"{name}_xi_angle.png", dpi=300)

if __name__ == '__main__':
    main()

